import cv2
import numpy as np
from rknnlite.api import RKNNLite

# 颜色常量定义
BLACK = (255, 255, 0)
WHITE = (255, 255, 255)
RED = (0, 0, 255)
GREEN = (0, 255, 0)
BLUE = (255, 0, 0)
GRAY = (128, 128, 128)
YELLOW = (0, 255, 255)

# RKNN模型配置
RKNN_MODEL = 'yolov5_qi.rknn'
OBJ_THRESH = 0.80
NMS_THRESH = 0
IMG_SIZE = 640
CLASSES = ("1", "2")  # 1代表黑棋，2代表白棋


class GomokuVision:
    def __init__(self, board_size=13, cell_size=36, start_x=130, start_y=5):
        self.board_size = board_size
        self.cell_size = cell_size
        self.start_x = start_x
        self.start_y = start_y
        self.rknn = self._init_rknn()
        self.cap = None
        self.detected_cells = [[False for _ in range(board_size)] for _ in range(board_size)]
        self.grid_drawn = False
        self.player_piece = 1  # 默认玩家执黑(1)

    def _init_rknn(self):
        """初始化RKNN模型"""
        rknn = RKNNLite()
        print('--> Load RKNN model')
        ret = rknn.load_rknn(RKNN_MODEL)
        if ret != 0:
            print('Load RKNN model failed')
            exit(ret)

        ret = rknn.init_runtime()
        if ret != 0:
            print('Init runtime environment failed!')
            exit(ret)

        print('RKNN model initialized successfully')
        return rknn

    def init_camera(self, camera_id=0):
        """初始化摄像头"""
        self.cap = cv2.VideoCapture(camera_id)
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

        if not self.cap.isOpened():
            print("Error: Could not open camera.")
            return False
        return True

    def set_player_piece(self, piece_type):
        """设置玩家棋子类型 (1=黑棋, 2=白棋)"""
        self.player_piece = piece_type

    def release(self):
        """释放资源"""
        if self.cap and self.cap.isOpened():
            self.cap.release()
        if self.rknn:
            self.rknn.release()
        cv2.destroyAllWindows()

    def xywh2xyxy(self, x):
        """将边界框从[x,y,w,h]格式转换为[x1,y1,x2,y2]格式"""
        y = np.copy(x)
        y[:, 0] = x[:, 0] - x[:, 2] / 2  # top left x
        y[:, 1] = x[:, 1] - x[:, 3] / 2  # top left y
        y[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right x
        y[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right y
        return y

    def process(self, input, mask, anchors):
        """处理模型输出，提取边界框、置信度和类别概率"""
        anchors = [anchors[i] for i in mask]
        grid_h, grid_w = map(int, input.shape[0:2])

        box_confidence = input[..., 4]
        box_confidence = np.expand_dims(box_confidence, axis=-1)

        box_class_probs = input[..., 5:]
        box_xy = input[..., :2] * 2 - 0.5

        col = np.tile(np.arange(0, grid_w), grid_w).reshape(-1, grid_w)
        row = np.tile(np.arange(0, grid_h).reshape(-1, 1), grid_h)
        col = col.reshape(grid_h, grid_w, 1, 1).repeat(3, axis=-2)
        row = row.reshape(grid_h, grid_w, 1, 1).repeat(3, axis=-2)
        grid = np.concatenate((col, row), axis=-1)

        box_xy += grid
        box_xy *= int(IMG_SIZE / grid_h)

        box_wh = pow(input[..., 2:4] * 2, 2)
        box_wh = box_wh * anchors

        box = np.concatenate((box_xy, box_wh), axis=-1)

        return box, box_confidence, box_class_probs

    def filter_boxes(self, boxes, box_confidences, box_class_probs):
        """过滤低置信度的边界框"""
        boxes = boxes.reshape(-1, 4)
        box_confidences = box_confidences.reshape(-1)
        box_class_probs = box_class_probs.reshape(-1, box_class_probs.shape[-1])

        _box_pos = np.where(box_confidences >= OBJ_THRESH)
        boxes = boxes[_box_pos]
        box_confidences = box_confidences[_box_pos]
        box_class_probs = box_class_probs[_box_pos]

        class_max_score = np.max(box_class_probs, axis=-1)
        classes = np.argmax(box_class_probs, axis=-1)
        _class_pos = np.where(class_max_score >= OBJ_THRESH)

        boxes = boxes[_class_pos]
        classes = classes[_class_pos]
        scores = (class_max_score * box_confidences)[_class_pos]

        return boxes, classes, scores

    def nms_boxes(self, boxes, scores):
        """非极大值抑制，去除重叠的边界框"""
        x = boxes[:, 0]
        y = boxes[:, 1]
        w = boxes[:, 2] - boxes[:, 0]
        h = boxes[:, 3] - boxes[:, 1]

        areas = w * h
        order = scores.argsort()[::-1]

        keep = []
        while order.size > 0:
            i = order[0]
            keep.append(i)

            xx1 = np.maximum(x[i], x[order[1:]])
            yy1 = np.maximum(y[i], y[order[1:]])
            xx2 = np.minimum(x[i] + w[i], x[order[1:]] + w[order[1:]])
            yy2 = np.minimum(y[i] + h[i], y[order[1:]] + h[order[1:]])

            w1 = np.maximum(0.0, xx2 - xx1 + 0.00001)
            h1 = np.maximum(0.0, yy2 - yy1 + 0.00001)
            inter = w1 * h1

            ovr = inter / (areas[i] + areas[order[1:]] - inter)
            inds = np.where(ovr <= NMS_THRESH)[0]
            order = order[inds + 1]

        keep = np.array(keep)
        return keep

    def yolov5_post_process(self, input_data):
        """YOLOv5模型的后处理"""
        masks = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
        anchors = [[46, 115], [48, 104], [51, 117], [53, 107], [56, 116], [59, 99],
                   [61, 124], [63, 109], [71, 100]]

        boxes, classes, scores = [], [], []
        for input, mask in zip(input_data, masks):
            b, c, s = self.process(input, mask, anchors)
            b, c, s = self.filter_boxes(b, c, s)
            boxes.append(b)
            classes.append(c)
            scores.append(s)

        boxes = np.concatenate(boxes)
        boxes = self.xywh2xyxy(boxes)
        classes = np.concatenate(classes)
        scores = np.concatenate(scores)

        nboxes, nclasses, nscores = [], [], []
        for c in set(classes):
            inds = np.where(classes == c)
            b = boxes[inds]
            c = classes[inds]
            s = scores[inds]

            keep = self.nms_boxes(b, s)

            nboxes.append(b[keep])
            nclasses.append(c[keep])
            nscores.append(s[keep])

        if not nclasses and not nscores:
            return None, None, None

        boxes = np.concatenate(nboxes)
        classes = np.concatenate(nclasses)
        scores = np.concatenate(nscores)

        return boxes, classes, scores

    def letterbox(self, im, new_shape=(640, 640), color=(0, 0, 0)):
        """调整图像大小并添加边界，保持原始宽高比"""
        shape = im.shape[:2]
        if isinstance(new_shape, int):
            new_shape = (new_shape, new_shape)

        r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])

        ratio = r, r
        new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
        dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]

        dw /= 2
        dh /= 2

        if shape[::-1] != new_unpad:
            im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
        top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
        left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
        im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)
        return im, ratio, (dw, dh)

    def draw_grid(self, image):
        """在图像上绘制网格"""
        # 绘制横线
        for i in range(self.board_size + 1):
            y = self.start_y + i * self.cell_size
            cv2.line(image, (self.start_x, y), (self.start_x + self.cell_size * self.board_size, y), BLACK, 1)

        # 绘制竖线
        for i in range(self.board_size + 1):
            x = self.start_x + i * self.cell_size
            cv2.line(image, (x, self.start_y), (x, self.start_y + self.cell_size * self.board_size), BLACK, 1)

        # 绘制坐标标签
        letters = [chr(ord('A') + i) for i in range(self.board_size)]

        # 上方坐标（字母表示列）
        for i in range(self.board_size):
            x = self.start_x + i * self.cell_size + self.cell_size // 2
            y = self.start_y - 10
            cv2.putText(image, letters[i], (x, y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, BLACK, 1)

        # 左侧坐标（数字表示行）
        for i in range(self.board_size):
            x = self.start_x - 15
            y = self.start_y + i * self.cell_size + self.cell_size // 2 + 5
            cv2.putText(image, str(i + 1), (x, y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, BLACK, 1)

        return image

    def get_grid_position(self, x, y):
        """获取坐标点所在的网格位置"""
        col = int((x - self.start_x) / self.cell_size)
        row = int((y - self.start_y) / self.cell_size)
        return row, col

    def fill_detected_cells(self, image):
        """用黄色填充已检测到的格子"""
        for row in range(self.board_size):
            for col in range(self.board_size):
                if self.detected_cells[row][col]:
                    x1 = self.start_x + col * self.cell_size
                    y1 = self.start_y + row * self.cell_size
                    x2 = x1 + self.cell_size
                    y2 = y1 + self.cell_size
                    cv2.rectangle(image, (x1, y1), (x2, y2), YELLOW, -1)
        return image

    def detect_chess_pieces(self, image):
        """使用RKNN模型检测棋子"""
        img, ratio, (dw, dh) = self.letterbox(image, new_shape=(IMG_SIZE, IMG_SIZE))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = np.expand_dims(img, 0)

        # 模型推理
        outputs = self.rknn.inference(inputs=[img])

        input0_data = outputs[0]
        input1_data = outputs[1]
        input2_data = outputs[2]

        input0_data = input0_data.reshape([3, -1] + list(input0_data.shape[-2:]))
        input1_data = input1_data.reshape([3, -1] + list(input1_data.shape[-2:]))
        input2_data = input2_data.reshape([3, -1] + list(input2_data.shape[-2:]))

        input_data = [np.transpose(input0_data, (2, 3, 0, 1)),
                      np.transpose(input1_data, (2, 3, 0, 1)),
                      np.transpose(input2_data, (2, 3, 0, 1))]

        boxes, classes, scores = self.yolov5_post_process(input_data)

        # 调整边界框坐标回原始图像尺寸
        if boxes is not None:
            boxes[:, [0, 2]] = (boxes[:, [0, 2]] - dw) / ratio[0]
            boxes[:, [1, 3]] = (boxes[:, [1, 3]] - dh) / ratio[1]

        return boxes, classes, scores

    def detect_player_move(self, frame):
        """检测玩家的落子位置"""
        # 先填充已检测的格子为黄色
        detect_frame = self.fill_detected_cells(frame.copy())

        # 检测棋子
        boxes, classes, scores = self.detect_chess_pieces(detect_frame)

        # 初始化棋盘状态 (0:空, 1:黑棋, 2:白棋)
        board_state = [[0 for _ in range(self.board_size)] for _ in range(self.board_size)]
        letters = [chr(ord('A') + i) for i in range(self.board_size)]
        player_move = None

        if boxes is not None:
            for box, cl in zip(boxes, classes):
                x1, y1, x2, y2 = box
                x_center = (x1 + x2) / 2
                y_center = (y1 + y2) / 2

                # 计算棋子中心所在的网格
                row, col = self.get_grid_position(x_center, y_center)

                # 检查网格坐标是否在有效范围内且未被检测过
                if 0 <= row < self.board_size and 0 <= col < self.board_size and not self.detected_cells[row][col]:
                    # 只检测玩家棋子类型
                    if cl + 1 == self.player_piece:
                        # 更新棋盘状态
                        board_state[row][col] = cl + 1

                        # 标记该格子已被检测
                        self.detected_cells[row][col] = True

                        # 记录玩家落子位置
                        player_move = (row, col)

                        # 获取网格中心坐标
                        grid_x = self.start_x + col * self.cell_size + self.cell_size // 2
                        grid_y = self.start_y + row * self.cell_size + self.cell_size // 2

                        # 绘制棋子标记
                        color = RED if cl == 0 else GREEN
                        cv2.circle(detect_frame, (int(grid_x), int(grid_y)), int(self.cell_size // 2 - 5), color, 2)

                        # 打印棋子位置
                        position = f"{letters[col]}{row + 1}"
                        piece_type = "B" if cl == 0 else "W"
                        print(f"Player {piece_type} {position}")

        # 绘制网格
        detect_frame = self.draw_grid(detect_frame)

        # 显示检测结果
        cv2.imshow("Player Move Detection", detect_frame)

        return player_move, board_state

    def reset_detection(self):
        """重置检测状态"""
        self.detected_cells = [[False for _ in range(self.board_size)] for _ in range(self.board_size)]
        print("Detection reset. All cells will be re-checked.")

    def toggle_grid(self):
        """切换网格显示状态"""
        self.grid_drawn = not self.grid_drawn
        print(f"Grid drawing {'started' if self.grid_drawn else 'stopped'}.")